"""
简单游戏演示 - 支持测试集黑方方案
提供一个基础的游戏循环和可视化界面，用于快速测试和演示

作者：wjttdbx
版本：4.0 - 集成测试集黑方方案
"""

import pygame
import sys
import os
import json
from typing import Dict, List, Tuple
import numpy as np

sys.path.insert(0, sys.path[0]+"/../")
from sh15.core.config import GameConfig
from sh15.agents import white_agent
from sh15.core.CoreEnvironment import CoreEnvironment
from sh15.agents.fsm_defense_agent import create_fsm_defense_agent
from sh15.agents.base_agent import AgentObservation
from sh15.agents.action_executor import ActionExecutor
from sh15.visualization import GameVisualizer
from sh15.core.LockingSystem import LockingSystem


class TestCaseLoader:
    """测试案例加载器"""
    
    def __init__(self, test_cases_dir: str):
        self.test_cases_dir = test_cases_dir
    
    def load_case(self, case_file: str) -> Dict:
        """加载单个测试案例"""
        file_path = os.path.join(self.test_cases_dir, case_file)
        
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        return {
            'case_name': case_file.replace('.json', ''),
            'black_boat_count': data.get('usvNum', 0),
            'black_boat_paths': {k: v for k, v in data.items() if k.startswith('usv') and k != 'usvNum'}
        }
    
    def load_all_cases(self) -> List[Dict]:
        """加载所有测试案例"""
        cases = []
        
        for filename in os.listdir(self.test_cases_dir):
            if filename.startswith("方案") and filename.endswith(".json"):
                case_data = self.load_case(filename)
                cases.append(case_data)
        
        # 按方案序号进行数字排序
        def get_case_number(case):
            try:
                # 从 '方案1-...' 中提取数字 1
                return int(case['case_name'].split('-')[0][2:])
            except (ValueError, IndexError):
                return -1  # 无法解析的名称排在前面
        
        cases.sort(key=get_case_number)
        return cases


class TestCaseBlackAgent:
    """测试案例专用黑方智能体"""
    
    def __init__(self, test_case_config: Dict):
        self.config = test_case_config
        self.boat_paths = {}
        self.boat_current_target = {}
        
    def reset(self):
        """重置智能体状态"""
        self.boat_paths = {}
        self.boat_current_target = {}
        
    def initialize_paths(self, black_boats: List):
        """初始化黑方无人艇的路径"""
        for i, boat in enumerate(black_boats):
            usv_key = f"usv{i+1}"
            if usv_key in self.config['black_boat_paths']:
                self.boat_paths[boat.id] = self.config['black_boat_paths'][usv_key]
                self.boat_current_target[boat.id] = 0
                
                # 设置初始位置为第一个路径点
                if self.boat_paths[boat.id]:
                    first_point = self.boat_paths[boat.id][0]["point"]
                    boat.x = first_point[0]
                    boat.y = first_point[1]
                    boat.speed = self.boat_paths[boat.id][0]["speed"]
    
    def update_all_black_boats(self, core_env, dt: float):
        """更新所有黑方无人艇：先按预设路径机动，再在不影响机动的前提下执行锁定尝试"""
        current_time = core_env.current_time
        for boat in core_env.black_boats:
            if not boat.active or boat.is_frozen(current_time):
                continue

            if boat.id not in self.boat_paths:
                continue

            # 先路径跟踪机动
            self._update_boat_path(boat, core_env.white_boats, dt)

            # 机动后刷新探测并尝试锁定（不改变机动，仅维护锁定状态机）
            LockingSystem.update_black_boat_detection(boat, core_env.white_boats)
            core_env.try_black_boat_lock(boat.id)
    
    def _update_boat_path(self, boat, white_boats: List, dt: float):
        """更新单艘船的路径跟踪"""
        path = self.boat_paths[boat.id]
        current_target_idx = self.boat_current_target[boat.id]
        
        if current_target_idx >= len(path):
            return
            
        target_point = path[current_target_idx]
        target_x, target_y = target_point["point"]
        target_speed = target_point["speed"]
        
        # 计算到目标点的距离
        dx = target_x - boat.x
        dy = target_y - boat.y
        distance = np.sqrt(dx**2 + dy**2)
        
        # 如果接近目标点，切换到下一个目标点
        if distance < target_speed * dt * 2:
            self.boat_current_target[boat.id] += 1
            if self.boat_current_target[boat.id] < len(path):
                next_target = path[self.boat_current_target[boat.id]]
                target_x, target_y = next_target["point"]
                target_speed = next_target["speed"]
                dx = target_x - boat.x
                dy = target_y - boat.y
                distance = np.sqrt(dx**2 + dy**2)
        
        # 移动向目标点
        if distance > 0:
            direction_x = dx / distance
            direction_y = dy / distance
            
            boat.speed = min(target_speed, GameConfig.BLACK_BOAT_MAX_SPEED)
            boat.heading = np.arctan2(direction_y, direction_x)
            
            move_distance = boat.speed * dt
            boat.x += direction_x * move_distance
            boat.y += direction_y * move_distance
            
        # 检查锁定
        self._attempt_locking(boat, white_boats)
    
    def _attempt_locking(self, boat, white_boats: List):
        """尝试锁定白方单位"""
        for white_boat in white_boats:
            if not white_boat.active:
                continue
                
            dx = white_boat.x - boat.x
            dy = white_boat.y - boat.y
            distance = np.sqrt(dx**2 + dy**2)
            
            if distance <= GameConfig.BLACK_BOAT_LOCK_RANGE:
                boat.target_id = white_boat.id
                boat.is_locking = True
                break


class SimpleGameDemo:
    """
    简单游戏演示类 - 支持测试集黑方方案
    
    提供一个基础的游戏循环和可视化界面
    """
    def __init__(self, width=1280, height=720, use_fsm_agent=False, test_cases_dir=None, headless: bool = False):
        """初始化游戏"""
        # 延后界面初始化到测试案例选择之后
        self.width = width
        self.height = height
        self.screen = None
        self.core_env = CoreEnvironment()
        
        # 选择智能体类型

        self.white_agent = create_fsm_defense_agent()
        self.agent_type = "FSM防御智能体"

            
        self.action_executor = ActionExecutor()
        
        # 启用调试模式以查看详细的动作执行信息
        self.action_executor.set_debug_mode(GameConfig.debug_mode)

        # 在选择测试案例后再创建可视化器
        self.visualization = None
        self.headless = headless
        
        self.running = False
        self.total_steps = 0
        
        # 测试集相关
        self.test_cases_dir = test_cases_dir
        self.test_case_loader = None
        self.available_cases = []
        self.selected_case = None
        self.black_agent = None
        
        if test_cases_dir and os.path.exists(test_cases_dir):
            self.test_case_loader = TestCaseLoader(test_cases_dir)
            self.available_cases = self.test_case_loader.load_all_cases()
            print(f"找到 {len(self.available_cases)} 个测试案例")

        # 日志相关
        self.log_data = {}
        self.next_log_time = 0.0
        self.log_file_path = None   # 聚合后的 JSON 文件路径，如 方案X_log.json
        self.log_jsonl_path = None  # 流式 JSONL 文件路径，如 方案X_log.jsonl
        self._log_jsonl_handle = None  # 持有文件句柄，减少频繁打开

    def select_test_case(self):
        """让用户选择测试案例"""
        if not self.available_cases:
            print("未找到测试案例，使用默认黑方策略")
            return False
            
        print("\n=== 可用的测试案例 ===")
        for i, case in enumerate(self.available_cases):
            print(f"{i+1}. {case['case_name']} (黑方艇数量: {case['black_boat_count']})")
        
        while True:
            try:
                choice = input(f"\n请选择测试案例 (1-{len(self.available_cases)}): ").strip()
                choice_num = int(choice)
                
                if 1 <= choice_num <= len(self.available_cases):
                    self.selected_case = self.available_cases[choice_num - 1]
                    print(f"已选择: {self.selected_case['case_name']}")
                    
                    # 设置黑方艇数量
                    GameConfig.BLACK_BOATS = self.selected_case['black_boat_count']
                    
                    # 创建测试案例黑方智能体
                    self.black_agent = TestCaseBlackAgent(self.selected_case)
                    self.black_agent_type = f"测试案例: {self.selected_case['case_name']}"
                    
                    # 确定日志文件路径（如 方案14_log.json），保存在 sh15/logs 目录
                    case_prefix = self.selected_case['case_name'].split('-')[0]
                    log_filename = f"{case_prefix}_log.json"
                    logs_dir = os.path.join('sh15', 'logs')
                    os.makedirs(logs_dir, exist_ok=True)
                    self.log_file_path = os.path.join(logs_dir, log_filename)
                    self.log_jsonl_path = os.path.join(logs_dir, f"{case_prefix}_log.jsonl")

                    # 交互式选择：是否禁用可视化（默认开启）
                    try:
                        ans = input("是否禁用可视化以加速运行? 输入 y 禁用，直接回车保持可视化 [y/N]: ").strip().lower()
                        if ans in ("y", "yes"):
                            self.headless = True
                    except KeyboardInterrupt:
                        print("\n用户取消输入，保持默认可视化模式")

                    return True
                else:
                    print(f"请输入 1 到 {len(self.available_cases)} 之间的数字")
                    
            except ValueError:
                print("请输入有效的数字")
            except KeyboardInterrupt:
                print("\n用户取消选择")
                return False

    def run_game_loop(self):
        """运行主游戏循环"""
        # 选择测试案例
        if not self.select_test_case():
            print("未选择测试案例，退出程序")
            return
        
        # 在选择测试案例之后再初始化界面与可视化器（headless时跳过）
        if not self.headless:
            pygame.init()
            self.screen = pygame.display.set_mode((self.width, self.height))
            pygame.display.set_caption("SH-15 博弈决策模型演示 - 测试集方案")
            self.visualization = GameVisualizer(self.screen)
        else:
            self.visualization = None
            
        self.running = True
        clock = pygame.time.Clock() if not self.headless else None
        
        # 打印策略说明
        print(f"=== SH-15 {self.agent_type} vs {self.black_agent_type} 演示 ===")
        print("策略说明:")
        print(f"- 黑方: 严格按照测试集中的路径点行进，遇到白方立即尝试锁定")
        print(f"- 白方: {self.agent_type}")
        print("==============================\n")
        
        # 重置环境和智能体
        self.core_env.reset()
        self.white_agent.reset()
        if self.black_agent:
            self.black_agent.reset()
        
        # 初始化黑方路径
        if self.black_agent:
            self.black_agent.initialize_paths(self.core_env.black_boats)
        
        print("游戏初始化完成")
        print(f"白方无人艇数量: {len(self.core_env.white_boats)}")
        print(f"黑方无人艇数量: {len(self.core_env.black_boats)}")
        print("开始游戏循环...\n")

        # 初始化日志器
        self.log_data = {}
        self.next_log_time = 0.0
        # 打开 JSONL 文件为写入模式（覆盖旧文件）
        try:
            if self.log_jsonl_path:
                self._log_jsonl_handle = open(self.log_jsonl_path, 'w', encoding='utf-8')
        except Exception as e:
            print(f"打开日志文件失败: {e}")
        
        while self.running:
            # 控制渲染帧率（仅可视化）
            if clock is not None:
                clock.tick(60)

            # 使用固定的游戏时间步长
            if self.visualization is not None:
                dt = GameConfig.TIME_STEP * self.visualization.speed_multiplier * 0.1
            else:
                dt = GameConfig.TIME_STEP

            self._handle_input()
            self._update_game_state(dt)
            self._render()
            self._check_game_over()
            
            self.total_steps += 1

    def _update_game_state(self, dt: float):
        """更新游戏状态"""
        # 1. 更新探测系统，让AI获取最新信息
        self.core_env.update_detection_system()
        
        # 2. AI决策
        # 创建观测对象
        # 仅传入“当前帧被探测到”的黑方目标，避免越权获取真实坐标
        detected_ids = set()
        for boat in self.core_env.white_boats:
            for t in boat.detected_targets:
                if t.active:
                    detected_ids.add(t.id)
        for drone in self.core_env.all_white_drones:
            for t in drone.detected_targets:
                if t.active:
                    detected_ids.add(t.id)

        detected_black_boats = [
            b for b in self.core_env.black_boats
            if b.active and b.id in detected_ids
        ]

        observation = AgentObservation(
            white_boats=self.core_env.white_boats,
            all_drones=self.core_env.all_white_drones,
            black_boats=detected_black_boats,
            current_time=self.core_env.current_time
        )
        
        agent_action = self.white_agent.get_actions(observation)
        if GameConfig.debug_mode:
            print(f"\n第{self.total_steps}步:")
            print(f"白方无人艇动作: {agent_action.boat_actions}")
            print(f"白方无人机动作: {agent_action.drone_actions}")
            print(f"当前时间: {self.core_env.current_time:.2f} 秒")
        
        # 3. 执行白方动作
        self.action_executor.execute_actions(
            agent_action,
            self.core_env.white_boats,
            self.core_env.all_white_drones,
            self.core_env,
            self.core_env.current_time,
            dt
        )
        
        # 4. 更新黑方单位（使用测试案例路径+锁定尝试）
        if self.black_agent:
            self.black_agent.update_all_black_boats(self.core_env, dt)

        # 5. 更新核心环境（物理、锁定、碰撞等）
        self.core_env.update(dt)

        # 6. 日志记录（每10秒一次）
        self._log_if_needed()

    def _handle_input(self):
        """处理用户输入"""
        if self.visualization is None:
            return
        for event in pygame.event.get():
            if not self.visualization.handle_event(event):
                self.running = False

    def _render(self):
        """渲染游戏画面"""
        if self.visualization is None:
            return
        game_state_info = self.core_env.get_game_state()
        
        # 获取状态机信息（统一构造，至少包含 predicted_black_targets）
        predicted_black_targets = {}
        try:
            detected_dict = getattr(self.white_agent, 'detected_targets', {}) or {}
            for unit_id, data in detected_dict.items():
                pos = data.get('predict_pos', data.get('position'))
                if pos is None or len(pos) < 2:
                    continue
                try:
                    key = int(unit_id)
                except Exception:
                    key = str(unit_id)
                predicted_black_targets[key] = (float(pos[0]), float(pos[1]))
        except Exception:
            predicted_black_targets = {}

        extras = {}
        if hasattr(self.white_agent, 'global_state') and hasattr(self.white_agent, 'boat_states'):
            task_manager = getattr(self.white_agent, 'task_manager', None)
            global_state = getattr(self.white_agent, 'global_state', None)
            extras = {
                'global_state': global_state.name if global_state else 'UNKNOWN',
                'state_timer': getattr(self.white_agent, 'state_timer', 0),
                'boat_states': {bid: state.name for bid, state in getattr(self.white_agent, 'boat_states', {}).items()},
                'drone_states': {did: state.name for did, state in getattr(self.white_agent, 'drone_states', {}).items()},
                'boat_targets': getattr(self.white_agent, 'boat_targets', {}).copy(),
                'task_assignment_ready': (task_manager is not None and 
                                          hasattr(task_manager, 'is_ready') and
                                          task_manager.is_ready()),
            }

        # 在交给可视化前，按当前活跃黑方过滤 predicted_black_targets，避免显示已击沉目标
        try:
            active_black_ids = {b.id for b in self.core_env.black_boats if b.active}
            filtered = {}
            for uid, pos in predicted_black_targets.items():
                try:
                    uid_int = int(uid)
                except Exception:
                    uid_int = None
                if (uid_int is not None and uid_int in active_black_ids) or (uid in active_black_ids) or (str(uid) in {str(x) for x in active_black_ids}):
                    filtered[uid] = pos
            predicted_black_targets = filtered
        except Exception:
            pass

        fsm_info = {'predicted_black_targets': predicted_black_targets, **extras}
        
        self.visualization.render_frame(
            white_boats=self.core_env.white_boats,
            white_drones=self.core_env.all_white_drones,
            black_boats=self.core_env.black_boats,
            current_time=self.core_env.current_time,
            game_info={ # 传递额外信息
                'collisions': self.core_env.collision_count,
                'intercepts': self.core_env.intercepted_black_boats,
                'fsm_info': fsm_info  # 新增：状态机信息
            }
        )
        pygame.display.flip()

    def _check_game_over(self):
        """检查游戏是否结束"""
        game_over, winner = self.core_env.is_game_over()
        if game_over:
            print(f"\n游戏结束: {winner}")
            self._show_game_over_screen(winner)
            self.running = False

    def _show_game_over_screen(self, winner: str):
        """显示游戏结束画面"""
        if self.visualization is None:
            # headless 模式下也打印关键统计
            game_stats = self.core_env.get_game_state()
            crossed_black = game_stats.get('crossed_black_boats', 0)
            collision_count = game_stats.get('collision_count', 0)
            white_lock_count = game_stats.get('white_lock_count', 0)
            black_lock_count = game_stats.get('black_lock_count', 0)
            first_break_time = game_stats.get('first_breakthrough_time', None)
            if first_break_time is None:
                intercept_time_text = "100小时"
            else:
                total_seconds = int(first_break_time)
                minutes = total_seconds // 60
                seconds = total_seconds % 60
                intercept_time_text = f"{minutes}分{seconds:02d}秒"

            print("\n=== 关键统计（控制台） ===")
            print(f"黑方突防成功数量: {crossed_black}")
            print(f"白方阻击时间: {intercept_time_text}")
            print(f"白方被锁定次数: {black_lock_count}")
            print(f"黑方被锁定次数: {white_lock_count}")
            print(f"白方发生碰撞次数: {collision_count}")
            print("======================\n")
            return
        game_stats = self.core_env.get_game_state()
        
        if winner == "white_wins":
            result_text = "白方胜利 (White Wins)"
        elif winner == "black_wins":
            result_text = "黑方胜利 (Black Wins)"
        else:
            result_text = "平局 (Draw)"

        # 控制台打印关键统计信息
        crossed_black = game_stats.get('crossed_black_boats', 0)
        collision_count = game_stats.get('collision_count', 0)
        white_lock_count = game_stats.get('white_lock_count', 0)
        black_lock_count = game_stats.get('black_lock_count', 0)
        first_break_time = game_stats.get('first_breakthrough_time', None)
        if first_break_time is None:
            intercept_time_text = "100小时"
        else:
            total_seconds = int(first_break_time)
            minutes = total_seconds // 60
            seconds = total_seconds % 60
            intercept_time_text = f"{minutes}分{seconds:02d}秒"

        print("\n=== 关键统计（控制台） ===")
        print(f"黑方突防成功数量: {crossed_black}")
        print(f"白方阻击时间: {intercept_time_text}")
        print(f"白方被锁定次数: {black_lock_count}")
        print(f"黑方被锁定次数: {white_lock_count}")
        print(f"白方发生碰撞次数: {collision_count}")
        print("======================\n")

        self._render()

        self.visualization.draw_game_over_overlay(
            result=result_text,
            game_stats={
                'steps': self.total_steps,
                'crossed_black_boats': game_stats.get('crossed_black_boats', 0),
                'eliminated_black_boats': game_stats.get('eliminated_black_boats', 0),
                'intercepted_black_boats': game_stats.get('intercepted_black_boats', 0),
                'collision_count': game_stats.get('collision_count', 0),
                'white_lock_count': game_stats.get('white_lock_count', 0),
                'black_lock_count': game_stats.get('black_lock_count', 0),
                'first_breakthrough_time': game_stats.get('first_breakthrough_time', None),
            },
            current_time=self.core_env.current_time
        )
        
        waiting_for_exit = True
        while waiting_for_exit:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    waiting_for_exit = False
                if event.type == pygame.KEYDOWN and (event.key == pygame.K_RETURN or event.key == pygame.K_ESCAPE):
                    waiting_for_exit = False
    
    def cleanup(self):
        """清理资源"""
        pygame.quit()
        print("资源清理完成")
        # 关闭 JSONL 句柄并转换为聚合 JSON
        try:
            if self._log_jsonl_handle is not None and not self._log_jsonl_handle.closed:
                self._log_jsonl_handle.close()
        finally:
            self._convert_jsonl_to_aggregated_json()

    # ================= 日志功能 =================
    def _log_if_needed(self):
        """当当前时间达到下一个10秒刻度时，记录一帧日志并写入文件"""
        if self.log_jsonl_path is None or self._log_jsonl_handle is None:
            return
        # 可能跨越多个10秒刻度，逐个补记
        while self.core_env.current_time >= self.next_log_time - 1e-6:
            timestamp_key = str(int(round(self.next_log_time)))
            snapshot = self._build_log_snapshot()
            line_obj = {"time": timestamp_key, **snapshot}
            try:
                json.dump(line_obj, self._log_jsonl_handle, ensure_ascii=False)
                self._log_jsonl_handle.write("\n")
                # 不强制 flush，每次换行写入由系统缓冲
            except Exception as e:
                print(f"写入日志失败: {e}")
            # 下一次记录时间
            self.next_log_time += 10.0

    def _build_log_snapshot(self) -> dict:
        """按要求构造一帧日志，结构参照 sh15/logs/logDemo.json"""
        from math import cos, sin

        def boat_state(boat) -> str:
            if not boat.active:
                return "退出"
            if boat.is_frozen(self.core_env.current_time):
                return "冻结"
            if getattr(boat, 'is_locking', False):
                return "锁定"
            return "移动"

        def drone_state(drone) -> str:
            if not drone.active:
                return "退出"
            # 停靠在母舰上
            if drone.docked_on_boat_id is not None and drone.docked_on_boat_id >= 0:
                # 能量充满视为“停泊”，否则为“补能”
                return "停泊" if drone.energy >= getattr(drone, 'max_energy', 100.0) else "补能"
            # 空中状态
            if drone.state in ("flying", "returning"):
                return "移动"
            # 兜底
            return "移动"

        def velocity_vector(unit) -> list:
            vx = unit.speed * cos(unit.heading)
            vy = unit.speed * sin(unit.heading)
            return [round(vx, 6), round(vy, 6), 0]

        def position_vector(unit) -> list:
            return [round(unit.x, 6), round(unit.y, 6), 0]

        white_section = {}
        # 白方无人艇
        for idx, boat in enumerate(self.core_env.white_boats, start=1):
            white_section[f"w_usv{idx}"] = {
                "pos": position_vector(boat),
                "type": "无人艇",
                "group": "白",
                "state": boat_state(boat),
                "velocity": velocity_vector(boat)
            }
        # 白方无人机（以列表顺序编号）
        for idx, drone in enumerate(self.core_env.all_white_drones, start=1):
            white_section[f"w_uav{idx}"] = {
                "pos": position_vector(drone),
                "type": "无人机",
                "group": "白",
                "state": drone_state(drone),
                "velocity": velocity_vector(drone)
            }

        black_section = {}
        for idx, boat in enumerate(self.core_env.black_boats, start=1):
            black_section[f"b_usv{idx}"] = {
                "pos": position_vector(boat),
                "type": "无人艇",
                "group": "黑",
                "state": boat_state(boat),
                "velocity": velocity_vector(boat)
            }

        return {
            "white": white_section,
            "black": black_section
        }

    def _convert_jsonl_to_aggregated_json(self):
        """将 JSONL 转换为聚合 JSON（键为时间串），写入 self.log_file_path"""
        if not self.log_jsonl_path or not self.log_file_path:
            return
        try:
            aggregated = {}
            with open(self.log_jsonl_path, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        obj = json.loads(line)
                        t = str(obj.get("time", ""))
                        if t:
                            # 保留 white/black 两节
                            aggregated[t] = {
                                "white": obj.get("white", {}),
                                "black": obj.get("black", {})
                            }
                    except Exception:
                        # 忽略单行解析错误
                        pass
            with open(self.log_file_path, 'w', encoding='utf-8') as out:
                json.dump(aggregated, out, ensure_ascii=False, indent=4)
        except Exception as e:
            print(f"转换日志失败: {e}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='SH-15博弈决策模型演示 - 测试集方案')
    parser.add_argument('--agent', choices=['simple', 'fsm'], default='simple',
                        help='选择白方智能体类型: simple(简单防线) 或 fsm(FSM状态机)')
    parser.add_argument('--test_dir', default='sh15/SH-15测试集',
                        help='测试案例目录路径')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    parser.add_argument('--headless', action='store_true', help='禁用可视化以加速运行')
    
    args = parser.parse_args()
    
    # 设置调试模式
    if args.debug:
        GameConfig.debug_mode = True
    
    use_fsm = (args.agent == 'fsm')
    demo = SimpleGameDemo(use_fsm_agent=use_fsm, test_cases_dir=args.test_dir, headless=args.headless)
    
    print(f"启动演示 - 白方: {demo.agent_type}")
    print("按 ESC 或关闭窗口退出")
    print("=" * 50)
    
    try:
        demo.run_game_loop()
    except (KeyboardInterrupt, SystemExit):
        print("\n用户中断游戏")
    except Exception as error:
        print(f"\n游戏出现错误: {error}")
        import traceback
        traceback.print_exc()
    finally:
        demo.cleanup()

